package timit;

import core.HiddenLayer;
import core.Phone;
import core.PhoneWrapper;
import core.Speaker;
import org.apache.commons.lang3.ArrayUtils;

import java.io.*;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Created with IntelliJ IDEA.
 * User: lahiru
 * Date: 6/11/14
 * Time: 4:30 PM
 * To change this template use File | Settings | File Templates.
 */
public class CalculatePhonemeStats {
    private static final double log2 = Math.log(2);
    private static int phoneCount = 117;
    private static final int numberOfHiddenUnits = 1024;
//    private static final String vowelIDs = "5 11 12 16 19 24 28 29 33 35 36 38 41 42";
//    private static final boolean writeFilesForClustering = false;
    private static final String haVectorDirectory = "/home/lahiru/experiments/Speaker-effects/vectors/train/";
    private static final String location = "/home/lahiru/experiments/Speaker-effects/objects/train/";

//    private static Map<Integer, Double[]> haVectors = new HashMap<Integer, Double[]>();

    private static Map<Integer, Map<Integer, Double>> h1PhoneDistances = new HashMap<Integer, Map<Integer, Double>>();        //Phoneme distances for each hidden layer
    private static Map<Integer, Map<Integer, Double>> h2PhoneDistances = new HashMap<Integer, Map<Integer, Double>>();
    private static Map<Integer, Map<Integer, Double>> h3PhoneDistances = new HashMap<Integer, Map<Integer, Double>>();
    private static Map<Integer, Map<Integer, Double>> h4PhoneDistances = new HashMap<Integer, Map<Integer, Double>>();
    private static Map<Integer, Map<Integer, Double>> h5PhoneDistances = new HashMap<Integer, Map<Integer, Double>>();


    private static double h1TotalPhoneDistance = 0;
    private static double h2TotalPhoneDistance = 0;
    private static double h3TotalPhoneDistance = 0;
    private static double h4TotalPhoneDistance = 0;
    private static double h5TotalPhoneDistance = 0;


    public static void main(String[] args) {

        try {
            System.out.println("Processing Hidden Layer 1");
            //read the hiddenlayer1
            FileInputStream fileIn = new FileInputStream(location + "h1.ser");
            ObjectInputStream in = new ObjectInputStream(fileIn);
            HiddenLayer hiddenLayer = (HiddenLayer) in.readObject();
            accumilatePhoneDifferences(hiddenLayer);

            //read the hiddenlayer2
            System.out.println("Processing Hidden Layer 2");
            fileIn = new FileInputStream(location + "h2.ser");
            in = new ObjectInputStream(fileIn);
            hiddenLayer = (HiddenLayer) in.readObject();
            accumilatePhoneDifferences(hiddenLayer);

            //read the hiddenlayer3
            System.out.println("Processing Hidden Layer 3");
            fileIn = new FileInputStream(location + "h3.ser");
            in = new ObjectInputStream(fileIn);
            hiddenLayer = (HiddenLayer) in.readObject();
            accumilatePhoneDifferences(hiddenLayer);

            //read the hiddenlayer4
            System.out.println("Processing Hidden Layer 4");
            fileIn = new FileInputStream(location + "h4.ser");
            in = new ObjectInputStream(fileIn);
            hiddenLayer = (HiddenLayer) in.readObject();
            accumilatePhoneDifferences(hiddenLayer);

            //read the hiddenlayer5
            System.out.println("Processing Hidden Layer 5");
            fileIn = new FileInputStream(location + "h5.ser");
            in = new ObjectInputStream(fileIn);
            hiddenLayer = (HiddenLayer) in.readObject();
            accumilatePhoneDifferences(hiddenLayer);


//            System.out.println("##########  Distances between HA Phonetic Vectors ##############");
            printPhoneHaDistances();


            in.close();
            fileIn.close();
        } catch (IOException i) {
            i.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        }

    }

    private static void accumilatePhoneDifferences(HiddenLayer hiddenLayer) throws IOException {

        Map<Integer, PhoneWrapper> phoneWrapperMap = new HashMap<Integer, PhoneWrapper>();

        for (Speaker speaker : hiddenLayer.getSpeakerList()) {
            Map<Integer, Phone> phones = speaker.getPhones();
            for (int i = 0; i < phones.size(); i++) {
                Phone phone;
                if (phones.get(i) == null) {
                    continue;
                }
                phone = phones.get(i);
                int phoneID = phone.getId();

                if (phoneWrapperMap.get(phoneID) == null) {
                    phoneWrapperMap.put(phoneID, new PhoneWrapper(phone));
                } else {
                    PhoneWrapper phoneWrapper = phoneWrapperMap.get(phoneID);
                    phoneWrapper.addActivations(phone.getVector());
                    phoneWrapper.increaseCount(phone.getCount());
                    phoneWrapperMap.put(phone.getId(), phoneWrapper);
                }
            }
        }

        Object[] tempPhoneIDs = phoneWrapperMap.keySet().toArray();
//        System.out.println("Before");
//        for(Object o : tempPhoneIDs){
//            System.out.println(o.toString());
//        }

        Arrays.sort(tempPhoneIDs);

//        System.out.println("After");
//        for(Object o : tempPhoneIDs){
//            System.out.println(o.toString());
//        }

        List<Object> phoneIdList = Arrays.asList(tempPhoneIDs);
        Map<Integer, Double[]> haVectors = new HashMap<Integer, Double[]>();

//        Average all the phonemeHiddenActivation Vectors
        for (int i = 0; i < phoneIdList.size(); i++) {
            int phoneID = (Integer) phoneIdList.get(i);
            PhoneWrapper phoneWrapper = phoneWrapperMap.get(phoneID);
            phoneWrapper.averageVector();
            phoneWrapperMap.put(phoneID, phoneWrapper);
            haVectors.put(phoneID, ArrayUtils.toObject(phoneWrapper.getVector()));
        }

        if(haVectors.size() != phoneCount){
            System.err.println("The number of vectors does not match the expected count");
            System.exit(1);
        }

        writeArkFile(hiddenLayer.getId(), haVectors);
      accumilatePhoneticHaDistances(hiddenLayer.getId(), phoneWrapperMap, phoneIdList);

    }

    private static void writeArkFile(int id, Map<Integer, Double[]> haVectors) throws IOException{
        FileWriter fileWriter = new FileWriter(haVectorDirectory + "pdfVectors_layer"+id + ".ark");
        BufferedWriter bufferedWriter = new BufferedWriter(fileWriter);

        bufferedWriter.write("PDF_LAYER_"+id + " ["+ "\n");

        for (int i = 0; i < phoneCount; i++) {

            StringBuffer sb = new StringBuffer();
            Double[] vector = haVectors.get(i);
            for (int p = 0; p < numberOfHiddenUnits; p++){
                if (p == numberOfHiddenUnits - 1) {
                    if(i==phoneCount-1){
                        sb.append(vector[p] + " ]"+ "\n");
                    }else{
                        sb.append(vector[p] + "\n");
                    }

                } else {
                    sb.append(vector[p] + " ");
                }
            }
            bufferedWriter.write(sb.toString());
        }
        bufferedWriter.close();
    }

    private static void accumilatePhoneticHaDistances(int hiddenLayerId, Map<Integer, PhoneWrapper> phoneWrapperMap, List<Object> phoneIdList) {
        double totalDistance = 0.0;

        for (int i = 0; i < phoneIdList.size(); i++) {

            Map<Integer, Double> distanceVector = new HashMap<Integer, Double>();

            int phoneID1 = (Integer) phoneIdList.get(i);

            PhoneWrapper phoneWrapper1 = phoneWrapperMap.get(phoneID1);
            for (int j = 0; j < phoneIdList.size(); j++) {
                int phoneID2 = (Integer) phoneIdList.get(j);
                PhoneWrapper phoneWrapper2 = phoneWrapperMap.get(phoneID2);

                //Calculate the distance between two phones
                double distance = calculateKLDivergenceVectors(phoneWrapper1.getVector(), phoneWrapper2.getVector());

                distanceVector.put(phoneID2, distance);
                totalDistance = totalDistance + distance;
            }
            if (hiddenLayerId == 1) {
                h1PhoneDistances.put(phoneID1, distanceVector);
            } else if (hiddenLayerId == 2) {
                h2PhoneDistances.put(phoneID1, distanceVector);
            } else if (hiddenLayerId == 3) {
                h3PhoneDistances.put(phoneID1, distanceVector);
            } else if (hiddenLayerId == 4) {
                h4PhoneDistances.put(phoneID1, distanceVector);
            } else if (hiddenLayerId == 5) {
                h5PhoneDistances.put(phoneID1, distanceVector);
            }
        }

        if (hiddenLayerId == 1) {
            h1TotalPhoneDistance = totalDistance;
        } else if (hiddenLayerId == 2) {
            h2TotalPhoneDistance = totalDistance;
        } else if (hiddenLayerId == 3) {
            h3TotalPhoneDistance = totalDistance;
        } else if (hiddenLayerId == 4) {
            h4TotalPhoneDistance = totalDistance;
        } else if (hiddenLayerId == 5) {
            h5TotalPhoneDistance = totalDistance;
        }
    }

    private static void printPhoneHaDistances() {
        System.out.println("Total core.Phone Distance for Hidden layer 1 " + h1TotalPhoneDistance);
        System.out.println("Average core.Phone Distance for Hidden layer 1 " + h1TotalPhoneDistance / phoneCount);

        System.out.println("Total core.Phone Distance for Hidden layer 2 " + h2TotalPhoneDistance);
        System.out.println("Average core.Phone Distance for Hidden layer 2 " + h2TotalPhoneDistance / phoneCount);

        System.out.println("Total core.Phone Distance for Hidden layer 3 " + h3TotalPhoneDistance);
        System.out.println("Average core.Phone Distance for Hidden layer 3 " + h3TotalPhoneDistance / phoneCount);

        System.out.println("Total core.Phone Distance for Hidden layer 4 " + h4TotalPhoneDistance);
        System.out.println("Average core.Phone Distance for Hidden layer 4 " + h4TotalPhoneDistance / phoneCount);

        System.out.println("Total core.Phone Distance for Hidden layer 5 " + h5TotalPhoneDistance);
        System.out.println("Average core.Phone Distance for Hidden layer 5 " + h5TotalPhoneDistance / phoneCount);
    }


    public static double calculateKLDivergenceVectors(double[] vector1, double[] vector2) throws NumberFormatException {
        int length = vector1.length;
        if (length != vector2.length) {
            throw new NumberFormatException("Vector length's are not matching ");
        }

        if (vector1.length != numberOfHiddenUnits) {
            System.out.println("vector lengths are wrong ");
        }

        double kLdiv = 0.0;
        for (int i = 0; i < length; i++) {

            String vector1s = vector1[i] + "";
            String vector2s = vector2[i] + "";

            if (vector1s.equals("NaN")) {
                System.out.println("The vector1 is NaN " + i);
            }

            if (vector2s.equals("NaN")) {
                System.out.println("The vector2 is NaN " + i);
            }

            double temp = klDivergence(vector1[i], vector2[i]);
            String temps = temp + "";
            if (temps.equals("NaN")) {
//                System.out.println("The temp is NaN " + i);
//                System.out.println(vector1[i]);
//                System.out.println(vector2[i]);
                temp = 0.0; //todo change
            }


//            kLdiv = kLdiv + klDivergence(vector1[i], vector2[i]);
            kLdiv = kLdiv + temp;
//            System.out.println("temp " + temp);
        }

        return kLdiv / length;  //To change body of created methods use File | Settings | File Templates.
    }

    public static double klDivergence(double p1, double p2) {


        double klDiv = 0.0;

        if (p1 == 0.0 || p2 == 0.0) {
            return klDiv;
        } else {
//            return ((p1 * Math.log(p1 / p2)) + ((1 - p1) * Math.log((1 - p1) / (1 - p2)))) / log2;
            return (((p1 * Math.log(p1 / p2)) + ((1 - p1) * Math.log((1 - p1) / (1 - p2)))) + ((p2 * Math.log(p2 / p1)) + ((1 - p2) * Math.log((1 - p2) / (1 - p1))))) / log2;
        }

    }
}
